package com.sicheng.datastruct.tree.huffman.huffmancode;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.ToString;

import java.util.*;

/**
 * @author zsc
 * @version 1.0
 * @date 2021/9/10 15:29
 */

@Data
@SuppressWarnings("unused")
public class HuffmanCode {

    public byte[] zip(String content) {
        Map<Byte, String> huffmanCode = getHuffmanCode(content);
        if (Objects.requireNonNull(huffmanCode).size() == 0)
            return null;
        byte[] contentBytes = content.getBytes();
        StringBuilder builder = new StringBuilder();
        for (byte contentByte : contentBytes) {
            builder.append(huffmanCode.get(contentByte));
        }
        //byte[] huffmanCodeBytes 的长度 len
        int len = (builder.length() + 7) / 8;

        int index = 0;
        byte[] huffmanCodeBytes = new byte[len];
        for (int i = 0; i < builder.length(); i += 8) {
            if (i + 8 > builder.length()) {
                huffmanCodeBytes[index++] = (byte) Integer.parseInt(builder.substring(i), 2);
            } else
                huffmanCodeBytes[index++] = (byte) Integer.parseInt(builder.substring(i, i + 8), 2);
        }
        return huffmanCodeBytes;
    }

    //HuffmanCode  使用Map来存储
    private Map<Byte, String> huffmanCode;

    public Map<Byte, String> getHuffmanCode(String content) {
        if (content == null || content.length() == 0)
            return null;
        Node root = initHuffmanTree(content);
        if (root == null)
            return null;
        huffmanCode = new HashMap<>();
        getHuffmanCode(root, "", new StringBuilder());
        return huffmanCode;
    }

    /**
     * @param root          huffman的根
     * @param code          决定向下的方向，  0向左孩子走，1向右孩子走
     * @param huffmanCoding 足步向下生成编码
     */
    private void getHuffmanCode(Node root, String code, StringBuilder huffmanCoding) {
        //向左，向右两个方向进行递归
        StringBuilder coding = new StringBuilder(huffmanCoding).append(code);
        if (root.getValue() == null) {
            getHuffmanCode(root.getLeft(), "0", coding);
            getHuffmanCode(root.getRight(), "1", coding);
        } else {
            huffmanCode.put(root.getValue(), coding.toString());
        }

    }

    //获取 Huffman树的根结点
    public static Node initHuffmanTree(String content) {

        if (content == null || content.length() == 0)
            return null;
        byte[] contentBytes = content.getBytes();

        List<Node> nodes = new ArrayList<>();
        //1.bytes 词频统计
        HashMap<Byte, Integer> count = new HashMap<>();
        for (byte v : contentBytes) {
            count.put(v, count.getOrDefault(v, 0) + 1);
        }
        count.forEach((key, value) -> nodes.add(new Node(key, value)));


        //构建哈夫曼树
        while (nodes.size() > 1) {
            nodes.sort(Comparator.comparingInt(e -> e.weight));

            Node left = nodes.get(0);
            Node right = nodes.get(1);

            Node parent = new Node(null, left.getWeight() + right.getWeight());

            parent.setLeft(left);
            parent.setRight(right);

            nodes.remove(left);
            nodes.remove(right);

            nodes.add(parent);
        }

        return nodes.get(0);
    }

    @Data
    @NoArgsConstructor
    @AllArgsConstructor
    @ToString(of = {"value", "weight"})
    static class Node {
        private Byte value;
        private int weight;
        private Node left;
        private Node right;

        public Node(Byte value, int weight) {
            this.value = value;
            this.weight = weight;
        }

        public void preOrder(Node root) {
            System.out.println(root);

            if (root.getLeft() != null)
                preOrder(root.getLeft());

            if (root.getRight() != null)
                preOrder(root.getRight());
        }
    }

}
